#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
FY-4A 全圆盘陆地 DEM 聚合 - 最终优化版

版本说明:
1.  读取 FY4A 原始坐标文件，并立即对 lat/lon 变量进行转置 (.T)
2.  创建一个全新的 NetCDF 文件，定义 ("y", "x") 维度。
3.  创建 1D (y, x) 和 2D (lat, lon) 坐标变量，以兼容 L2 数据。
4.  创建 2D "elev_mean" 等数据变量。
5.  [已修正] Worker 函数能正确处理 ds.index() 返回的 (row, col) 元组。
6.  [已优化] Main 函数使用批处理 (BATCH_SIZE) 提交任务，以防止完整
    运行时内存溢出。
7.  [已设置] ENABLE_TEST_FILTER 设置为 False，以运行完整全球数据。
"""
import os
import glob
import math
import numpy as np
import rasterio
from rtree import index as rtree_index
from concurrent.futures import ProcessPoolExecutor, as_completed
from tqdm import tqdm
import netCDF4 as nc
from global_land_mask import globe
import multiprocessing
import warnings
import logging

# --- 配置 ---
warnings.filterwarnings("ignore", category=RuntimeWarning)
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")

# -------------------------
# 用户参数（请按需修改）
# -------------------------
FY4A_NC = "/home/liudd/data_preprocessing/FY4A_coordinates.nc"
DEM_ROOT = "/mnt/raid1/liudd/DEM_land"
OUT_NC = "/home/liudd/data_preprocessing/FY4A_DEM_L2_COMPATIBLE_FULL.nc"

# 半径（km）: 2 km 半径 -> 4km 方形窗口
RADIUS_KM = 2.0
# 根据您的 CPU 核心数设置
N_WORKERS = max(1, int(multiprocessing.cpu_count() * 0.75))
# 批处理大小，防止内存溢出
BATCH_SIZE = 5000

# === 运行模式设置 ===
# 设置为 False 以运行完整的全球陆地数据
ENABLE_TEST_FILTER = False

# (测试区域边界，仅在 ENABLE_TEST_FILTER = True 时激活)
TEST_LAT_MIN = 25.0
TEST_LAT_MAX = 45.0
TEST_LON_MIN = 70.0
TEST_LON_MAX = 105.0


# -------------------------


# -------------------------
# 辅助：扫描 DEM 瓦片并建立空间索引
# -------------------------
def build_dem_index(dem_root):
    """
    扫描 dem_root 下所有 .tif 文件，读取元数据并构建 R-tree 空间索引。
    """
    logging.info(f"扫描 DEM 瓦片目录: {dem_root} ...")
    pattern = os.path.join(dem_root, "**", "*.tif")
    files = glob.glob(pattern, recursive=True)
    logging.info(f"找到 {len(files)} 个 .tif 文件")

    props = []
    idx = rtree_index.Index()
    for i, f in enumerate(files):
        try:
            with rasterio.open(f) as src:
                b = src.bounds
                t = src.transform
                props.append({
                    "id": i, "path": f,
                    "left": b.left, "right": b.right, "bottom": b.bottom, "top": b.top,
                    "transform": t, "width": src.width, "height": src.height,
                    "nodatavals": src.nodatavals  # 存储 nodata 值
                })
                # rtree 插入 (left, bottom, right, top)
                idx.insert(i, (b.left, b.bottom, b.right, b.top))
        except Exception as e:
            logging.warning(f"无法读取瓦片 {f}，跳过。错误: {e}")

    logging.info(f"成功索引 {len(props)} 个可读瓦片")
    return props, idx


# -------------------------
# 辅助：计算点窗口 bbox（经纬度）
# -------------------------
def point_window_bbox(lon, lat, radius_km=RADIUS_KM):
    """
    返回矩形窗口经纬度范围 (minx, miny, maxx, maxy)
    """
    lat_radius = radius_km / 111.0
    cos_lat = math.cos(math.radians(lat))
    if abs(cos_lat) < 1e-6:
        cos_lat = 1e-6
    lon_radius = radius_km / (111.0 * cos_lat)
    minx = lon - lon_radius
    maxx = lon + lon_radius
    miny = lat - lat_radius
    maxy = lat + lat_radius
    return minx, miny, maxx, maxy


# -------------------------
# Worker：(已修正 Bug)
# -------------------------
def worker_process_point(args):
    """
    args = (idx_flat, row_idx, col_idx, lon, lat, dem_meta_list)
    返回 (idx_flat, row_idx, col_idx, mean, std, min, max, n_valid)
    """
    # DEM_OPEN_CACHE 是一个进程内的全局缓存
    global DEM_OPEN_CACHE
    if "DEM_OPEN_CACHE" not in globals():
        DEM_OPEN_CACHE = {}

    idx_flat, row_i, col_j, lon, lat, dem_meta_list = args
    minx, miny, maxx, maxy = point_window_bbox(lon, lat, RADIUS_KM)

    values = []
    n_valid_total = 0
    for meta in dem_meta_list:
        # 1. 瓦片相交测试 (R-tree 已在 main 中完成，这里是冗余检查)
        # if meta["right"] < minx or meta["left"] > maxx or meta["top"] < miny or meta["bottom"] > maxy:
        #     continue
        path = meta["path"]

        try:
            # 2. 打开/复用瓦片
            ds = DEM_OPEN_CACHE.get(path)
            if ds is None:
                ds = rasterio.open(path)
                DEM_OPEN_CACHE[path] = ds

            # 3. 计算精确的像素窗口
            inter_minx = max(minx, meta["left"])
            inter_maxx = min(maxx, meta["right"])
            inter_miny = max(miny, meta["bottom"])
            inter_maxy = min(maxy, meta["top"])

            # === 关键修正：正确分配 ds.index() 返回的 (row, col) ===
            # ds.index(x, y) 返回 (row, col)
            row_LL, col_LL = ds.index(inter_minx, inter_miny)  # 左下角 (Lower-Left)
            row_UR, col_UR = ds.index(inter_maxx, inter_maxy)  # 右上角 (Upper-Right)

            # 使用正确的变量进行 min/max
            r0 = min(row_LL, row_UR)
            r1 = max(row_LL, row_UR)
            c0 = min(col_LL, col_UR)
            c1 = max(col_LL, col_UR)
            # ========================================================

            # 规范化并钳制到瓦片边界
            r0 = max(0, r0)
            c0 = max(0, c0)
            r1 = min(ds.height - 1, r1)
            c1 = min(ds.width - 1, c1)

            if r1 < r0 or c1 < c0:
                continue

            # 4. 读取数据
            window = ((r0, r1 + 1), (c0, c1 + 1))
            arr = ds.read(1, window=window, boundless=True, fill_value=np.nan)
            arr = arr.astype(float)

            # 5. 处理 NoData
            nodata = meta.get("nodatavals")
            if nodata is not None and nodata[0] is not None:
                # 检查 nodata 是否为 NaN，如果是，则无需额外处理
                if not np.isnan(nodata[0]):
                    arr[arr == nodata[0]] = np.nan

            vals = arr[~np.isnan(arr)]
            if vals.size > 0:
                values.append(vals)
                n_valid_total += vals.size
        except Exception as e:
            # logging.warning(f"Worker 瓦片处理出错 {path}: {e}")
            continue

    if n_valid_total == 0:
        return (idx_flat, row_i, col_j, np.nan, np.nan, np.nan, np.nan, 0)

    try:
        merged = np.concatenate(values)
    except Exception:  # Fallback for memory issue
        merged = np.hstack([v for v in values if v.size > 0])

    if merged.size == 0:
        return (idx_flat, row_i, col_j, np.nan, np.nan, np.nan, np.nan, 0)

    # 6. 计算统计
    mean_elev = float(np.mean(merged))
    std_elev = float(np.std(merged))
    min_elev = float(np.min(merged))
    max_elev = float(np.max(merged))

    return (idx_flat, row_i, col_j, mean_elev, std_elev, min_elev, max_elev, n_valid_total)


# -------------------------
# 主流程 (已修正结构 + 批处理优化)
# -------------------------
def main():
    # 1) 读取 FY4A netCDF 的 lat/lon (必须转置)
    logging.info("打开 FY4A netCDF，读取并转置经纬度 ...")
    try:
        with nc.Dataset(FY4A_NC, "r") as src_nc:
            # === 关键修正 1：必须转置 (.T) ===
            lat_2d = np.array(src_nc.variables['lat'][:, :]).T
            lon_2d = np.array(src_nc.variables['lon'][:, :]).T
    except Exception as e:
        logging.error(f"无法读取或转置坐标文件 {FY4A_NC}: {e}")
        return

    # 处理 NaN 值
    lat_2d[np.isnan(lat_2d)] = -90.0
    lon_2d[np.isnan(lon_2d)] = 360.0

    # shape 现在是转置后的 (Y, X)
    shape = lat_2d.shape
    n_rows, n_cols = shape  # n_rows 是 "y" 维度, n_cols 是 "x" 维度
    logging.info(f"FY4A 网格尺寸 (转置后 Y, X): {n_rows} x {n_cols}")

    # 2) 扫描并索引 DEM 瓦片
    dem_meta_list, dem_rtree = build_dem_index(DEM_ROOT)
    if len(dem_meta_list) == 0:
        logging.error("没有找到任何可读 DEM 瓦片，程序退出。")
        return

    # 3) === 关键修正 2：创建全新的 NetCDF 输出文件 ===
    logging.info(f"创建 *全新* L2 兼容 NetCDF: {OUT_NC}")
    # 使用 with 语句确保文件在出错时也能正确关闭
    with nc.Dataset(OUT_NC, "w", format="NETCDF4") as dst_nc:

        # A. 创建 ("y", "x") 维度
        dst_nc.createDimension("y", n_rows)
        dst_nc.createDimension("x", n_cols)

        # B. 创建 1D 坐标变量 y 和 x
        y_coord = dst_nc.createVariable("y", "i4", ("y",))
        x_coord = dst_nc.createVariable("x", "i4", ("x",))
        y_coord[:] = np.arange(n_rows)
        x_coord[:] = np.arange(n_cols)
        y_coord.long_name = "y index (row)"
        x_coord.long_name = "x index (column)"
        y_coord.units = "1"
        x_coord.units = "1"

        # C. 创建 2D lat/lon 变量
        dims_2d = ("y", "x")
        out_lat = dst_nc.createVariable("lat", "f4", dims_2d)
        out_lon = dst_nc.createVariable("lon", "f4", dims_2d)
        out_lat.long_name = "latitude"
        out_lon.long_name = "longitude"
        out_lat.units = "degrees_north"
        out_lon.units = "degrees_east"

        logging.info("正在将转置后的 (Y,X) 经纬度写入新文件...")
        out_lat[:] = lat_2d
        out_lon[:] = lon_2d

        # D. 创建高程变量
        elev_mean = dst_nc.createVariable("elev_mean", "f4", dims_2d, zlib=True, complevel=4, fill_value=np.nan)
        elev_std = dst_nc.createVariable("elev_std", "f4", dims_2d, zlib=True, complevel=4, fill_value=np.nan)
        elev_min = dst_nc.createVariable("elev_min", "f4", dims_2d, zlib=True, complevel=4, fill_value=np.nan)
        elev_max = dst_nc.createVariable("elev_max", "f4", dims_2d, zlib=True, complevel=4, fill_value=np.nan)
        elev_n = dst_nc.createVariable("elev_nvalid", "i4", dims_2d, zlib=True, complevel=4, fill_value=-9999)

        # 4) 准备待处理点列表
        logging.info("筛选出陆地区域点索引...")
        if ENABLE_TEST_FILTER:
            logging.warning("!!! 快速测试模式已启用 !!!")
        else:
            logging.info("--- 完整运行模式已启用 ---")

        land_indices = []
        flat_idx = 0

        # 预加载数据到内存以加快循环速度
        lat_2d_local = lat_2d[:]
        lon_2d_local = lon_2d[:]

        for i in range(n_rows):  # i 是 "y" 索引 (row)
            for j in range(n_cols):  # j 是 "x" 索引 (col)
                lat_ij = float(lat_2d_local[i, j])
                lon_ij = float(lon_2d_local[i, j])

                if lat_ij == -90.0 or lon_ij == 360.0:  # 跳过 NaN 点
                    flat_idx += 1
                    continue

                lon_check = lon_ij
                if lon_check > 180:
                    lon_check = lon_check - 360

                # === 区域筛选 ===
                if ENABLE_TEST_FILTER:
                    if not (TEST_LAT_MIN <= lat_ij <= TEST_LAT_MAX and
                            TEST_LON_MIN <= lon_check <= TEST_LON_MAX):
                        flat_idx += 1
                        continue

                try:
                    if globe.is_land(lat_ij, lon_check):
                        land_indices.append((flat_idx, i, j, lon_ij, lat_ij))
                except Exception:
                    pass
                flat_idx += 1

        logging.info(f"需要处理的陆地点数: {len(land_indices)}")

        # 5) 准备并行计算任务
        logging.info("准备并行计算任务...")
        tasks = []
        for (flat_idx, i, j, lon_ij, lat_ij) in tqdm(land_indices, desc="Building task list"):
            minx, miny, maxx, maxy = point_window_bbox(lon_ij, lat_ij, RADIUS_KM)
            # 使用 R-tree 索引快速查找相交的瓦片 ID
            candidate_ids = list(dem_rtree.intersection((minx, miny, maxx, maxy)))
            # 仅传递相交瓦片的元数据
            candidate_metas = [dem_meta_list[cid] for cid in candidate_ids] if candidate_ids else []

            # 如果没有候选瓦片，但该点是陆地（例如，小岛屿），也提交任务
            # worker 将会返回 n_valid=0 和 NaN
            tasks.append((flat_idx, i, j, lon_ij, lat_ij, candidate_metas))

        # 6) === 关键优化 3：批处理执行任务 ===
        logging.info(f"开始并行计算，worker 数: {N_WORKERS}, 批次大小: {BATCH_SIZE}")

        with ProcessPoolExecutor(max_workers=N_WORKERS) as executor:

            # 按 BATCH_SIZE 逐批提交任务
            for batch_start in range(0, len(tasks), BATCH_SIZE):
                batch_end = min(batch_start + BATCH_SIZE, len(tasks))
                batch_tasks = tasks[batch_start:batch_end]

                futures = [executor.submit(worker_process_point, t) for t in batch_tasks]

                batch_desc = f"Processing Batch {batch_start // BATCH_SIZE + 1}/{math.ceil(len(tasks) / BATCH_SIZE)}"

                # 处理当前批次返回的结果
                for fut in tqdm(as_completed(futures), total=len(futures), desc=batch_desc):
                    try:
                        res = fut.result()
                    except Exception as e:
                        logging.warning(f"Worker 出错: {e}")
                        continue

                    # res: (flat_idx, row_i, col_j, mean, std, min, max, n_valid)
                    _, ri, cj, mean_v, std_v, min_v, max_v, n_v = res

                    # 7) 写入 NetCDF (使用 [ri, cj] 直接写入)
                    try:
                        elev_mean[ri, cj] = mean_v
                        elev_std[ri, cj] = std_v
                        elev_min[ri, cj] = min_v
                        elev_max[ri, cj] = max_v
                        elev_n[ri, cj] = int(n_v)
                    except Exception as e:
                        logging.warning(f"写 netCDF 时出错 at (y={ri}, x={cj}): {e}")

    # 8) 结束
    logging.info(f"所有完成！最终输出已保存到: {OUT_NC}")


if __name__ == "__main__":
    main()